import os
import nn
import tvm
import time
import argparse
import mxnet as mx
import nnvm.compiler
from tvm import autotvm
from tvm.contrib import graph_runtime as runtime

def compile(args):
    if not os.path.exists(args.tvm):
        os.makedirs(args.tvm)
    shape_dict = {'data': args.shape}
    x = mx.nd.random.randn(*shape_dict['data'])
    sym, arg_params, aux_params = mx.model.load_checkpoint(args.mxnet_prefix, 0)
    sym, params = nnvm.frontend.from_mxnet(sym, arg_params, aux_params)
    with autotvm.apply_history_best(args.log):
        with nnvm.compiler.build_config(opt_level=3):
            graph, lib, params = nnvm.compiler.build(sym, 'llvm -mcpu=core-avx2', shape_dict, 'float32', params)
    lib.export_library(args.tvm_prefix + '.so')
    with open(args.tvm_prefix + '.json', 'w') as f:
        f.write(graph.json())
    with open(args.tvm_prefix + '.params', 'wb') as f:
        f.write(nnvm.compiler.save_param_dict(params))

def raw_speed(args):
    if not os.path.exists(args.mxnet):
        os.makedirs(args.mxnet)
    shape_dict = {'data': args.shape}
    x = mx.nd.random.randn(*shape_dict['data'])
    net = nn.get_net(args.network)
    net.initialize()
    # params = 'models/mtcnn/{}/{:0>3}.params'.format(args.network, args.epoch)
    # net.load_parameters(params, ctx=mx.cpu())
    net.hybridize(), net(x), net.export(args.mxnet_prefix, 0)
    x = x.asnumpy()
    sym, arg_params, aux_params = mx.model.load_checkpoint(args.mxnet_prefix, 0)
    sym, params = nnvm.frontend.from_mxnet(sym, arg_params, aux_params)
    with nnvm.compiler.build_config(opt_level=3):
        graph, lib, params = nnvm.compiler.build(sym, 'llvm -mcpu=core-avx2', shape_dict, 'float32', params)
    mod = runtime.create(graph, lib, tvm.cpu())
    mod.load_params(bytearray(nnvm.compiler.save_param_dict(params)))
    time.sleep(2)
    t0 = time.time()
    for _ in range(100):
        mod.run(data=x)
    print('raw speed {}: {:.3f}'.format(args.network, 10 * (time.time() - t0)))

def tune_speed(args):
    graph = open(args.tvm_prefix + '.json').read()
    lib = tvm.module.load(args.tvm_prefix + '.so')
    params = bytearray(open(args.tvm_prefix + '.params', 'rb').read())
    x = mx.nd.random.randn(*args.shape).asnumpy()
    module = runtime.create(graph, lib, tvm.cpu())
    module.load_params(params)
    time.sleep(2)
    t0 = time.time()
    for _ in range(100):
        module.run(data=x)
    print('tune speed {}: {:.3f}'.format(args.network, (time.time() - t0) * 10))

def parse_args():
    parser = argparse.ArgumentParser(description='Compile tvm model.')
    parser.add_argument('network', type=str)
    parser.add_argument('epoch', type=int)
    parser.add_argument('size', type=str, help='12,12')
    parser.add_argument('--tvm', type=str, default="models/tvm")
    parser.add_argument('--mxnet', type=str, default="models/mxnet")
    args = parser.parse_args()
    args.shape = [1,3] + list(map(int, args.size.split(',')))
    args.log = 'models/autotvm/{}.{}.log'.format(args.network, args.size)
    args.tvm_prefix = os.path.join(args.tvm, '.'.join((args.network, args.size)))
    args.mxnet_prefix = os.path.join(args.mxnet, args.network)
    return args

if __name__ == '__main__':
    args = parse_args()
    raw_speed(args)
    compile(args)
    tune_speed(args)

